{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from copy import deepcopy\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import pandas as pd\n",
    "import torch.nn as nn\n",
    "import random\n",
    "from tqdm import tqdm\n",
    "from sklearn.model_selection import train_test_split\n",
    "from torch.utils.data import Dataset, DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "#设置种子\n",
    "def set_seed(seed = 427):\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    os.environ['PYTHONHASHSEED'] = str(seed)\n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "#数据导入\n",
    "def load_data2():\n",
    "    # CMIP data    \n",
    "    train = xr.open_dataset('tcdata/enso_round1_train_20210201/CMIP_train.nc')\n",
    "    label = xr.open_dataset('tcdata/enso_round1_train_20210201/CMIP_label.nc')    \n",
    "   \n",
    "    train_sst = train['sst'][:, :12].values  # (4645, 12, 24, 72)截取前12项\n",
    "    train_t300 = train['t300'][:, :12].values\n",
    "    train_ua = train['ua'][:, :12].values\n",
    "    train_va = train['va'][:, :12].values\n",
    "    train_label = label['nino'][:, 12:36].values\n",
    "    \n",
    "    train_ua = np.nan_to_num(train_ua)#缺失值补0\n",
    "    train_va = np.nan_to_num(train_va)\n",
    "    train_t300 = np.nan_to_num(train_t300)\n",
    "    train_sst = np.nan_to_num(train_sst)\n",
    "\n",
    "    # SODA data    \n",
    "    train2 = xr.open_dataset('tcdata/enso_round1_train_20210201/SODA_train.nc')\n",
    "    label2 = xr.open_dataset('tcdata/enso_round1_train_20210201/SODA_label.nc')\n",
    "    train_sst2 = train2['sst'][:, :12].values  # (4645, 12, 24, 72)\n",
    "    train_t3002 = train2['t300'][:, :12].values\n",
    "    train_ua2 = train2['ua'][:, :12].values\n",
    "    train_va2 = train2['va'][:, :12].values\n",
    "    train_label2 = label2['nino'][:, 12:36].values\n",
    "    dict_train = {\n",
    "        'sst':train_sst,\n",
    "        't300':train_t300,\n",
    "        'ua':train_ua,\n",
    "        'va': train_va,\n",
    "        'label': train_label}\n",
    "    dict_valid = {\n",
    "        'sst':train_sst2,\n",
    "        't300':train_t3002,\n",
    "        'ua':train_ua2,\n",
    "        'va': train_va2,\n",
    "        'label': train_label2}\n",
    "    train_dataset = EarthDataSet(dict_train)\n",
    "    valid_dataset = EarthDataSet(dict_valid)\n",
    "    return train_dataset, valid_dataset\n",
    "    \n",
    "class EarthDataSet(Dataset):\n",
    "    def __init__(self, data):\n",
    "        self.data = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data['sst'])\n",
    "\n",
    "    def __getitem__(self, idx):   \n",
    "        return (self.data['sst'][idx], self.data['t300'][idx], self.data['ua'][idx], self.data['va'][idx]), self.data['label'][idx]\n",
    "    \n",
    "\n",
    "#模型\n",
    "class simpleSpatailTimeNN(nn.Module):\n",
    "    def __init__(self, n_cnn_layer:int=1, kernals:list=[3], n_lstm_units:int=64):\n",
    "        super(simpleSpatailTimeNN, self).__init__()\n",
    "        self.conv1 = nn.ModuleList([nn.Conv2d(in_channels=12, out_channels=12, kernel_size=i) for i in kernals]) \n",
    "        self.conv2 = nn.ModuleList([nn.Conv2d(in_channels=12, out_channels=12, kernel_size=i) for i in kernals])\n",
    "        self.conv3 = nn.ModuleList([nn.Conv2d(in_channels=12, out_channels=12, kernel_size=i) for i in kernals])\n",
    "        self.conv4 = nn.ModuleList([nn.Conv2d(in_channels=12, out_channels=12, kernel_size=i) for i in kernals])\n",
    "        self.pool1 = nn.AdaptiveAvgPool2d((22, 1))\n",
    "        self.pool2 = nn.AdaptiveAvgPool2d((1, 70))\n",
    "        self.pool3 = nn.AdaptiveAvgPool2d((1, 128))\n",
    "        self.batch_norm = nn.BatchNorm1d(12, affine=False)\n",
    "        self.lstm = nn.LSTM(1540 * 4, n_lstm_units, 2, bidirectional=True)\n",
    "        self.linear = nn.Linear(128, 24)\n",
    "\n",
    "    def forward(self, sst, t300, ua, va):\n",
    "        for conv1 in self.conv1:\n",
    "            sst = conv1(sst)  # batch * 12 * (24 - 2) * (72 -2)\n",
    "        for conv2 in self.conv2:\n",
    "            t300 = conv2(t300)\n",
    "        for conv3 in self.conv3:\n",
    "            ua = conv3(ua)\n",
    "        for conv4 in self.conv4:\n",
    "            va = conv4(va)\n",
    "\n",
    "        sst = torch.flatten(sst, start_dim=2)  # batch * 12 * 1540\n",
    "        t300 = torch.flatten(t300, start_dim=2)\n",
    "        ua = torch.flatten(ua, start_dim=2)\n",
    "        va = torch.flatten(va, start_dim=2)  # if flat, lstm input_dims = 1540 * 4              \n",
    "            \n",
    "        x = torch.cat([sst, t300, ua, va], dim=-1) #在内层合并 batch*12*1540*4\n",
    "        x = self.batch_norm(x)\n",
    "        x, _ = self.lstm(x)#输入1540*4,64hidden,2layer且双向 输出batch*(128=64*2)\n",
    "        x = self.pool3(x).squeeze(dim=-2)\n",
    "        x = self.linear(x)#128-->24\n",
    "        return x\n",
    "\n",
    "#loss function\n",
    "def coreff(x, y):\n",
    "    x_mean = np.mean(x)\n",
    "    y_mean = np.mean(y)\n",
    "    c1 = sum((x - x_mean) * (y - y_mean))\n",
    "    c2 = sum((x - x_mean)**2) * sum((y - y_mean)**2)\n",
    "    return c1/np.sqrt(c2)\n",
    "\n",
    "def rmse(preds, y):\n",
    "    return np.sqrt(sum((preds - y)**2)/preds.shape[0])\n",
    "\n",
    "def eval_score(preds, label):\n",
    "    # preds = preds.cpu().detach().numpy().squeeze()\n",
    "    # label = label.cpu().detach().numpy().squeeze()\n",
    "    acskill = 0\n",
    "    RMSE = 0\n",
    "    a = 0\n",
    "    a = [1.5]*4 + [2]*7 + [3]*7 + [4]*6\n",
    "    for i in range(24):\n",
    "        RMSE += rmse(label[:, i], preds[:, i])\n",
    "        cor = coreff(label[:, i], preds[:, i])\n",
    "    \n",
    "        acskill += a[i] * np.log(i+1) * cor\n",
    "    return 2/3 * acskill - RMSE\n",
    "\n",
    "#train\n",
    "fit_params = {\n",
    "    'n_epochs' : 22,\n",
    "    'learning_rate' : 8e-5,\n",
    "    'batch_size' : 64,\n",
    "}\n",
    "\n",
    "#下面的和原来有更改\n",
    "def train(model):\n",
    "    set_seed()\n",
    "    train_dataset, valid_dataset = load_data2()      \n",
    "    train_loader = DataLoader(train_dataset, batch_size=fit_params['batch_size'])\n",
    "    valid_loader = DataLoader(valid_dataset, batch_size=fit_params['batch_size'])\n",
    "    #model = simpleSpatailTimeNN()\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'   \n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=fit_params['learning_rate'])\n",
    "    loss_fn = nn.MSELoss()   \n",
    "    \n",
    "    model.to(device)\n",
    "    loss_fn.to(device)\n",
    "\n",
    "    for i in range(fit_params['n_epochs']):\n",
    "        model.train()\n",
    "        for step, ((sst, t300, ua, va), label) in enumerate(train_loader):                \n",
    "            sst = sst.to(device).float()\n",
    "            t300 = t300.to(device).float()\n",
    "            ua = ua.to(device).float()\n",
    "            va = va.to(device).float()\n",
    "            optimizer.zero_grad()\n",
    "            label = label.to(device).float()\n",
    "            preds = model(sst, t300, ua, va)\n",
    "            loss = loss_fn(preds, label)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            print('Step: {}, Train Loss: {}'.format(step, loss))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model= simpleSpatailTimeNN()\n",
    "model_trained=train(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=np.load('tcdata/enso_round1_test_20210201/test_0144-01-12.npy')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_test(x):\n",
    "    device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "    test_sst = x[:,:,:,0] \n",
    "    test_sst=torch.tensor(test_sst.reshape(1,12,24,72))\n",
    "    test_t300 = x[:,:,:,1] \n",
    "    test_t300=torch.tensor(test_t300.reshape(1,12,24,72))\n",
    "    test_ua = x[:,:,:,3] \n",
    "    test_ua=torch.tensor(test_ua.reshape(1,12,24,72))\n",
    "    test_va = x[:,:,:,0] \n",
    "    test_va=torch.tensor(test_va.reshape(1,12,24,72))\n",
    "    sst = test_sst.to(device).float()\n",
    "    t300 = test_t300.to(device).float()\n",
    "    ua = test_ua.to(device).float()\n",
    "    va = test_va.to(device).float()\n",
    "    preds = model_trained(sst, t300, ua, va)\n",
    "    return preds.view(-1).cuda().data.cpu().numpy()#要写这些把GPU上tensor转到CPU numpy上"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 0.05039866,  0.05358052,  0.01949791,  0.11609137,  0.09350426,\n",
       "        0.02987252,  0.02812253,  0.06287099,  0.05091866,  0.04479775,\n",
       "        0.0383049 ,  0.11089746, -0.01401671,  0.02064119,  0.05917842,\n",
       "        0.02108178,  0.02627509,  0.10765471,  0.06344109,  0.00503268,\n",
       "       -0.02946471,  0.07113085,  0.03924801,  0.01660974], dtype=float32)"
      ]
     },
     "execution_count": 110,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "predict_test(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Predict from /home/lvzhihao/JupyterLab/Tianchi-AIEarth-2021\n",
      "tcdata/enso_round1_test_20210201/test_0144-01-12.npy\n"
     ]
    }
   ],
   "source": [
    "test_dir = 'tcdata/enso_round1_test_20210201'\n",
    "print('Predict from ' + os.getcwd())\n",
    "for i in os.listdir(test_dir):\n",
    "    if i.endswith('.npy'):\n",
    "        x = np.load(test_dir + '/' + i)\n",
    "        np.save('result/' + i, predict_test(x))\n",
    "        print(test_dir + '/' + i)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
